'''
通过重叠区域的融合处理，合并标注中的重叠框
1. 重叠区域定义： 两个框的交集面积与最小框面积的比值大于阈值
2. 合并方式： 取两个框的最小外包矩形作为新的框 

'''
import os
import json
import shutil
import random
from collections import defaultdict
from PIL import Image, ImageDraw, ImageFont

# ==== 配置 ====
data_dir = '/home/guoyi/Dataset/aier_processed'  
ori_annotation_name = 'annotations.json'
save_orignal_annotation_name = 'annotations_v1.json'
tar_annotation_name = 'annotations.json'
image_dir = os.path.join(data_dir, 'images')
experiments_dir = './experiments'
visual_dir = os.path.join(experiments_dir, 'visual_overlap_merge')
os.makedirs(visual_dir, exist_ok=True)
os.system(f'rm -rf {visual_dir}/*')# 删除之前的可视化目录

font_path = './SarasaMonoCL-Regular.ttf'

threshold = 0.5
sample_n = 30

# ==== 1. 备份标注文件 ====
ori_path = os.path.join(data_dir, ori_annotation_name)
bak_path = os.path.join(data_dir, save_orignal_annotation_name)
tar_path = os.path.join(data_dir, tar_annotation_name)
if not os.path.exists(bak_path):
    shutil.copy(ori_path, bak_path)
    print(f"已备份原始标注到: {bak_path}")
else:
    print("备份已存在，跳过")

with open(ori_path, 'r', encoding='utf-8') as f:
    annotations = json.load(f)

def overlap_min_area(b1, b2):
    # 输入: [xmin, ymin, xmax, ymax], 归一化
    x1 = max(b1[0], b2[0])
    y1 = max(b1[1], b2[1])
    x2 = min(b1[2], b2[2])
    y2 = min(b1[3], b2[3])
    inter_w = max(0, x2 - x1)
    inter_h = max(0, y2 - y1)
    inter_area = inter_w * inter_h
    area1 = (b1[2] - b1[0]) * (b1[3] - b1[1])
    area2 = (b2[2] - b2[0]) * (b2[3] - b2[1])
    min_area = min(area1, area2)
    if min_area == 0:
        return 0
    return inter_area / min_area
def merge_bbox(b1, b2):
    # 返回两个框的并集（最小外包矩形）
    x1 = min(b1[0], b2[0])
    y1 = min(b1[1], b2[1])
    x2 = max(b1[2], b2[2])
    y2 = max(b1[3], b2[3])
    return [x1, y1, x2, y2]

# ==== 2. 重叠融合处理 ====
annotations_new = {}
merge_info_images = []  # 记录哪些图片被merge过
for img_name, info in annotations.items():
    lesions = info['lesions']
    # 分name分组
    name2bbox_idx = defaultdict(list)
    for idx, lesion in enumerate(lesions):
        name2bbox_idx[lesion['name']].append(idx)
    to_remove = set()
    merged_lesions = []
    merged_flag = False
    used_idx = set()
    for name, idxs in name2bbox_idx.items():
        idxs = [i for i in idxs if i not in used_idx]
        n = len(idxs)
        merged = [False]*n
        # 两两判断
        for i in range(n):
            if merged[i]: continue
            bbox_i = lesions[idxs[i]]['bbox']
            group = [idxs[i]]
            for j in range(i+1, n):
                if merged[j]: continue
                bbox_j = lesions[idxs[j]]['bbox']
                overlap_ratio = overlap_min_area(bbox_i, bbox_j)
                if overlap_ratio > threshold:
                    # 合并，交集最大
                    bbox_new = merge_bbox(bbox_i, bbox_j)
                    if bbox_new is not None:
                        merged_lesions.append({
                            "name": name,
                            "bbox": bbox_new
                        })
                        merged[i] = merged[j] = True
                        used_idx.add(idxs[i])
                        used_idx.add(idxs[j])
                        merged_flag = True
            if not merged[i] and idxs[i] not in used_idx:
                # 保留原来的
                merged_lesions.append(lesions[idxs[i]])
                used_idx.add(idxs[i])
    # 可能有未处理的（未merge也未遍历到）
    for idx in range(len(lesions)):
        if idx not in used_idx:
            merged_lesions.append(lesions[idx])
    if merged_flag:
        merge_info_images.append(img_name)
    # 其它信息拷贝
    info_new = dict(info)
    info_new['lesions'] = merged_lesions
    annotations_new[img_name] = info_new

# ==== 3. 保存新标注 ====
with open(tar_path, 'w', encoding='utf-8') as f:
    json.dump(annotations_new, f, ensure_ascii=False, indent=2)
print(f"已保存merge重叠后的标注到: {tar_path}")

# ==== 4. 抽样可视化重叠前后（before/after） ====
font = ImageFont.truetype(font_path, 16)
if len(merge_info_images) > 0:
    sample_images = random.sample(merge_info_images, min(len(merge_info_images), sample_n))
    for img_name in sample_images:
        img_path = os.path.join(image_dir, img_name)
        try:
            img = Image.open(img_path).convert("RGB")
            # before
            img_before = img.copy()
            draw = ImageDraw.Draw(img_before)
            for lesion in annotations[img_name]['lesions']:
                bbox = lesion['bbox']
                x0 = int(bbox[0] * 224)
                y0 = int(bbox[1] * 224)
                x1 = int(bbox[2] * 224)
                y1 = int(bbox[3] * 224)
                draw.rectangle([x0, y0, x1, y1], outline="red", width=2)
                draw.text((x0, y0), lesion['name'], font=font, fill="yellow")
            img_before.save(os.path.join(visual_dir, img_name.replace('.jpg', '_before.jpg')))

            # after
            img_after = img.copy()
            draw = ImageDraw.Draw(img_after)
            for lesion in annotations_new[img_name]['lesions']:
                bbox = lesion['bbox']
                x0 = int(bbox[0] * 224)
                y0 = int(bbox[1] * 224)
                x1 = int(bbox[2] * 224)
                y1 = int(bbox[3] * 224)
                draw.rectangle([x0, y0, x1, y1], outline="green", width=2)
                draw.text((x0, y0), lesion['name'], font=font, fill="yellow")
            img_after.save(os.path.join(visual_dir, img_name.replace('.jpg', '_after.jpg')))
        except Exception as e:
            print(f"可视化失败: {img_name}, 错误: {e}")

print(f"已采样{len(merge_info_images)}张合并图片，采样保存前后对比图（最多{sample_n}张）到 {visual_dir}")

